#include "GetFilesInDir.h"
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <strsafe.h>

#pragma comment (lib, "strsafe.lib")
#pragma comment (lib, "User32.lib")

void WINAPI_HandleError(LPTSTR lpszFunction) 
{ 
	// Retrieve the system error message for the last-error code

	LPVOID lpMsgBuf;
	LPVOID lpDisplayBuf;
	DWORD dw = GetLastError(); 

	FormatMessage(
		FORMAT_MESSAGE_ALLOCATE_BUFFER | 
		FORMAT_MESSAGE_FROM_SYSTEM |
		FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL,
		dw,
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		(LPTSTR) &lpMsgBuf,
		0, NULL );

	// Display the error message and exit the process
	lpDisplayBuf = (LPVOID)LocalAlloc(LMEM_ZEROINIT, 
		(lstrlen((LPCTSTR)lpMsgBuf)+lstrlen((LPCTSTR)lpszFunction)+40)*sizeof(char)); 
	StringCchPrintf((LPTSTR)lpDisplayBuf, 
		LocalSize(lpDisplayBuf) / sizeof(char),
		TEXT("%s failed with error %d: %s"), 
		lpszFunction, dw, lpMsgBuf);
	MessageBox(NULL, (LPCTSTR)lpDisplayBuf, TEXT("Error"), MB_OK); 

	LocalFree(lpMsgBuf);
	LocalFree(lpDisplayBuf);
}

void GetFilesInDir(const std::string& dir, std::vector<std::string>& files)
{
	WIN32_FIND_DATA ffd;
	HANDLE hFind = INVALID_HANDLE_VALUE;
	DWORD dwError = 0;

	std::string cppStrDir(dir);
	if (cppStrDir.at(cppStrDir.size() - 1) == '\\')
		cppStrDir.erase(cppStrDir.size() - 1);

	// Find the first file in the directory.
	hFind = FindFirstFile((cppStrDir + "\\*").c_str(), &ffd);

	if (INVALID_HANDLE_VALUE == hFind) 
	{
		WINAPI_HandleError(TEXT("FindFirstFile"));
		return ;
	} 

	// Push all the files in the dir
	do
	{
		if (!(ffd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY))
			files.push_back(ffd.cFileName);
	}
	while (FindNextFile(hFind, &ffd) != 0);

	dwError = GetLastError();

	if (dwError != ERROR_NO_MORE_FILES) 
		WINAPI_HandleError(TEXT("FindFirstFile"));

	FindClose(hFind);
}